//  Copyright (c) Microsoft Corporation.  All Rights Reserved.

using System;
using System.Collections.Generic;
using System.Text;
using System.Xml;
using System.Globalization;
using System.IO;
using System.Security.Cryptography.Xml;
using System.Collections.ObjectModel;

using System.Runtime.Serialization;

using System.IdentityModel.Selectors;
using System.IdentityModel.Tokens;

namespace EAI.SAML
{
    public class EAISamlAssertion
    {
        string version = "SAML 2.0";
        string assertionID = "SAML_" + Guid.NewGuid().ToString();
        DateTime issueInstant = DateTime.UtcNow.ToUniversalTime();

        // saml:Issuer element.
        string nameQualifier;
        string spNameQualifier;
        string issuerFormat;
        string spProviderID;
        string issuer;

        SecurityToken signingToken = null;
        SecurityKey signingKey = null;
        MemoryStream assertionRead = null;

        EAISamlSubject subject;
        EAISamlConditions conditions;
        EAISamlAdvice advice;

        readonly ImmutableCollection<EAISamlStatement> statements = new ImmutableCollection<EAISamlStatement>();

        SigningCredentials signingCredentials;

        ReadOnlyCollection<SecurityKey> cryptoList;

        bool isReadOnly = false;

        public EAISamlAssertion()
        {
        }

        public EAISamlAssertion(XmlReader reader, EAISamlSerializer serializer, SecurityTokenSerializer keyInfoSerializer, SecurityTokenResolver outOfBandTokenResolver)
        {
            this.ReadFrom(reader, serializer, keyInfoSerializer, outOfBandTokenResolver);
        }

        public void MakeReadOnly()
        {
            if (!this.isReadOnly)
            {
                this.statements.MakeReadOnly();
                if (this.advice != null)
                    this.advice.MakeReadOnly();
                if (this.conditions != null)
                    this.conditions.MakeReadOnly();
                if (this.subject != null)
                    this.subject.MakeReadOnly();
                if (this.cryptoList == null)
                    this.cryptoList = BuildCryptoList();

                this.isReadOnly = true;
            }
        }

        public string IssuerNameQualifier
        {
            get { return this.nameQualifier; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                this.nameQualifier = value;
            }
        }

        public string IssuerSPNameQualifier
        {
            get { return this.spNameQualifier; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                this.spNameQualifier = value;
            }
        }

        public string IssuerNameFormat
        {
            get { return this.issuerFormat; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                this.issuerFormat = value;
            }
        }

        public string IssuerSPProviderID
        {
            get { return this.spProviderID; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                this.spProviderID = value;
            }
        }

        public string Issuer
        {
            get { return this.issuer; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                if (String.IsNullOrEmpty(value))
                    throw new ArgumentException("Issuer cannot be null or empty string.");

                this.issuer = value;
            }
        }

        public string AssertionID
        {
            get { return this.assertionID; }
            set
            {
                if (String.IsNullOrEmpty(assertionID))
                    throw new ArgumentException("assertionID");

                this.assertionID = value;
            }
        }

        public EAISamlSubject Subject
        {
            get { return this.subject; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                if (value == null)
                    throw new ArgumentNullException("value");

                this.subject = value;
            }
        }

        public EAISamlConditions Conditions
        {
            get { return this.conditions; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                if (value == null)
                    throw new ArgumentNullException("value");

                this.conditions = value;
            }
        }

        public EAISamlAdvice Advice
        {
            get { return this.advice; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The object is read-only.");

                if (value == null)
                    throw new ArgumentNullException("value");

                this.advice = value;
            }
        }

        public IList<EAISamlStatement> Statements
        {
            get
            {
                return this.statements;
            }
        }

        public SigningCredentials SigningCredentials
        {
            get { return this.signingCredentials; }
            set
            {
                if (isReadOnly)
                    throw new InvalidOperationException("The objec is read-only.");

                this.signingCredentials = value;
            }
        }

        internal ReadOnlyCollection<SecurityKey> CryptoList
        {
            get { return this.cryptoList; }
        }

        internal SecurityToken SigningToken
        {
            get { return this.signingToken; }
        }

        ReadOnlyCollection<SecurityKey> BuildCryptoList()
        {
            List<SecurityKey> cryptos = new List<SecurityKey>();
            if (this.Subject != null)
            {
                foreach (EAISamlSubjectConfirmation subjectConfirmation in this.Subject.SubjectConfirmations)
                {
                    foreach (SecurityKey securityKey in subjectConfirmation.SubjectConfirmationData.SecurityKeys)
                        cryptos.Add(securityKey);
                }
            }

            return cryptos.AsReadOnly();
        }

        void CheckObjectValidity()
        {
            if (this.Issuer == null)
                throw new InvalidOperationException("saml:Assertion requires the issuer to be specified.");
        }

        public virtual void ReadFrom(XmlReader reader, EAISamlSerializer serializer, SecurityTokenSerializer keyInfoSerializer, SecurityTokenResolver outOfBandTokenResolver)
        {
            this.assertionRead = new MemoryStream();
            XmlDictionaryWriter writer = XmlDictionaryWriter.CreateTextWriter(this.assertionRead);
            writer.WriteNode(reader, false);
            writer.Flush();
            this.assertionRead.Position = 0;

            XmlDocument dom = new XmlDocument();
            dom.Load(this.assertionRead);

            this.assertionRead.Position = 0;
            XmlDictionaryReader dicReader = XmlDictionaryReader.CreateTextReader(this.assertionRead.ToArray(), XmlDictionaryReaderQuotas.Max);
            ReadXml(dicReader, serializer, keyInfoSerializer, outOfBandTokenResolver);

            SecurityKeyIdentifier keyIdentifier = null;

            XmlNode sigNode = dom.DocumentElement.ChildNodes[1];
            if ((sigNode.LocalName == "Signature") && (sigNode.NamespaceURI == "http://www.w3.org/2000/09/xmldsig#"))
            {
                dom.DocumentElement.RemoveChild(sigNode);
                for (int i = 0; i < sigNode.ChildNodes.Count; ++i)
                {
                    if ((sigNode.ChildNodes[i].LocalName == "KeyInfo") && (sigNode.ChildNodes[i].NamespaceURI == "http://www.w3.org/2000/09/xmldsig#"))
                    {
                        XmlDictionaryReader keyIdentifierReader = XmlDictionaryReader.CreateTextReader(Encoding.UTF8.GetBytes(sigNode.ChildNodes[i].OuterXml), XmlDictionaryReaderQuotas.Max);
                        keyIdentifier = EAISamlSerializer.ReadSecurityKeyIdentifier(keyIdentifierReader, keyInfoSerializer);
                        this.signingToken = EAISamlSerializer.ResolveSecurityToken(keyIdentifier, outOfBandTokenResolver);
                        this.signingKey = EAISamlSerializer.ResolveSecurityKey(keyIdentifier, outOfBandTokenResolver);
                        sigNode.RemoveChild(sigNode.ChildNodes[i]);
                        break;
                    }
                }
            }

            SignedXml signature = new SignedXml(dom);
            signature.LoadXml((XmlElement)sigNode);
            bool sigVerified = false;
            if (this.signingKey is AsymmetricSecurityKey)
            {
                sigVerified = signature.CheckSignature(((AsymmetricSecurityKey)this.signingKey).GetAsymmetricAlgorithm(signature.SignatureMethod, false));
            }
            else
            {
                sigVerified = signature.CheckSignature(((SymmetricSecurityKey)this.signingKey).GetKeyedHashAlgorithm(signature.SignatureMethod));
            }

            if (!sigVerified)
                throw new SecurityTokenException("Unable to validate the Signature on SAML.");
        }

        public virtual void ReadXml(XmlReader reader, EAISamlSerializer serializer, SecurityTokenSerializer keyInfoSerializer, SecurityTokenResolver outOfBandTokenResolver)
        {
            if (reader == null)
                throw new ArgumentNullException("reader");

            if (serializer == null)
                throw new ArgumentNullException("serializer");

            if (!reader.IsStartElement(EAISamlConstants.Assertion, EAISamlConstants.Namespace))
                throw new ArgumentException("The reader is not at <saml:Assertion> element.");

            this.version = reader.GetAttribute(EAISamlConstants.VersionAttribute);
            if (String.IsNullOrEmpty(this.version))
                throw new SecurityTokenException(String.Format("<saml:Assertion> is missing the required {0} attribute.", EAISamlConstants.VersionAttribute));

            this.assertionID = reader.GetAttribute(EAISamlConstants.IDAttribute);
            if (String.IsNullOrEmpty(this.assertionID))
                throw new SecurityTokenException(String.Format("<saml:Assertion> is missing the required {0} attribute", EAISamlConstants.IDAttribute));

            string attrValue = reader.GetAttribute(EAISamlConstants.IssueInstantAttribute);
            if (String.IsNullOrEmpty(attrValue))
                this.issueInstant = DateTime.ParseExact(attrValue, EAISamlConstants.AcceptedDateTimeFormats, DateTimeFormatInfo.InvariantInfo, DateTimeStyles.None).ToUniversalTime();

            reader.Read();

            // Read saml:Issuer element.
            reader.ReadStartElement(EAISamlConstants.Issuer, EAISamlConstants.Namespace);
            this.nameQualifier = reader.GetAttribute(EAISamlConstants.NameQualifierAttribute);
            this.spNameQualifier = reader.GetAttribute(EAISamlConstants.SPNameQualifierAttribute);
            this.spProviderID = reader.GetAttribute(EAISamlConstants.SPProviderIDAttribute);
            this.issuerFormat = reader.GetAttribute(EAISamlConstants.FormatAttribute);
            this.issuer = reader.ReadString();
            reader.ReadEndElement();

            if (reader.IsStartElement("Signature", "http://www.w3.org/2000/09/xmldsig#"))
            {
                reader.Skip();
            }

            // Read saml:Subject element, if present.
            if (reader.IsStartElement(EAISamlConstants.Subject, EAISamlConstants.Namespace))
                this.subject = serializer.LoadSubject(reader, keyInfoSerializer, outOfBandTokenResolver);

            // Read saml:Conditions element, if present.
            if (reader.IsStartElement(EAISamlConstants.Conditions))
                this.conditions = serializer.LoadConditions(reader, keyInfoSerializer, outOfBandTokenResolver);

            // Read saml:Advice element, if present.
            if (reader.IsStartElement(EAISamlConstants.Advice))
                this.advice = serializer.LoadAdvice(reader, keyInfoSerializer, outOfBandTokenResolver);

            // Read all the saml:Statements.
            while (reader.IsStartElement())
            {
                this.statements.Add(serializer.LoadStatement(reader, keyInfoSerializer, outOfBandTokenResolver));
            }

            if ((this.statements == null) || (this.statements.Count == 0))
                throw new SecurityTokenException("saml:Assertion should have at least one saml:Statement element.");

            reader.ReadEndElement();
        }

        public virtual void WriteTo(XmlWriter writer, SecurityTokenSerializer tokenSerializer)
        {
            if (this.assertionRead != null)
            {
                // We have already read this assertion. Just serialize 
                // this out again.
                XmlDictionaryReader reader = XmlDictionaryReader.CreateTextReader(this.assertionRead.GetBuffer(), XmlDictionaryReaderQuotas.Max);
                writer.WriteNode(reader, false);
                return;
            }

            if (this.signingCredentials == null)
                throw new InvalidOperationException("No Signing Credentials specified. Cannot serialize a unsigned SAML token.");

            MemoryStream ms = new MemoryStream();
            XmlDictionaryWriter dicWriter = XmlDictionaryWriter.CreateTextWriter(ms);
            this.WriteXml(dicWriter, tokenSerializer);
            dicWriter.Flush();
            ms.Position = 0;

            XmlDocument dom = new XmlDocument();
            dom.Load(ms);
            XmlNode assertionNode = dom.DocumentElement;

            XmlElement signatureElement = SignAssertion(dom, tokenSerializer);
            XmlNode firstChild = assertionNode.ChildNodes[0];
            assertionNode.InsertAfter(signatureElement, firstChild);

            assertionNode.WriteTo(writer);
        }

        public virtual void WriteXml(XmlWriter writer, SecurityTokenSerializer tokenSerializer)
        {
            CheckObjectValidity();

            if (writer == null)
                throw new ArgumentNullException("writer");

            writer.WriteStartElement(EAISamlConstants.Assertion, EAISamlConstants.Namespace);
            writer.WriteAttributeString(EAISamlConstants.VersionAttribute, this.version);
            writer.WriteAttributeString(EAISamlConstants.IDAttribute, this.assertionID);
            writer.WriteAttributeString(EAISamlConstants.IssueInstantAttribute, this.issueInstant.ToString(EAISamlConstants.GeneratedDateTimeFormat, DateTimeFormatInfo.InvariantInfo));

            // saml:Issuer
            writer.WriteStartElement(EAISamlConstants.Issuer, EAISamlConstants.Namespace);
            if (!String.IsNullOrEmpty(this.nameQualifier))
                writer.WriteAttributeString(EAISamlConstants.NameQualifierAttribute, this.nameQualifier);
            if (!String.IsNullOrEmpty(this.spNameQualifier))
                writer.WriteAttributeString(EAISamlConstants.SPNameQualifierAttribute, this.spNameQualifier);
            if (!String.IsNullOrEmpty(this.spProviderID))
                writer.WriteAttributeString(EAISamlConstants.SPProviderIDAttribute, this.spProviderID);
            if (!String.IsNullOrEmpty(this.issuerFormat))
                writer.WriteAttributeString(EAISamlConstants.FormatAttribute, this.issuerFormat);
            writer.WriteString(this.issuer);
            writer.WriteEndElement();

            // saml:Subject
            if (this.subject != null)
                this.subject.WriteXml(writer, tokenSerializer);

            // saml:Conditions
            if (this.conditions != null)
                this.conditions.WriteXml(writer, tokenSerializer);

            // saml:Advice
            if (this.advice != null)
                this.advice.WriteXml(writer, tokenSerializer);

            for (int i = 0; i < this.statements.Count; ++i)
                this.statements[i].WriteXml(writer, tokenSerializer);

            writer.WriteEndElement();
        }

        // Signs the Assertion that is passed in as a dom using the SigningCredentials
        // specified.
        XmlElement SignAssertion(XmlDocument dataToSign, SecurityTokenSerializer tokenSerializer)
        {
            SignedXml signedXml = new SignedXml(dataToSign);

            Reference reference = new Reference();
            reference.Uri = "#" + this.assertionID;
            reference.AddTransform(new XmlDsigC14NTransform());
            reference.AddTransform(new XmlDsigEnvelopedSignatureTransform());

            signedXml.AddReference(reference);

            if (this.signingCredentials.SigningKey is AsymmetricSecurityKey)
            {
                signedXml.SigningKey = ((AsymmetricSecurityKey)this.signingCredentials.SigningKey).GetAsymmetricAlgorithm(this.signingCredentials.SignatureAlgorithm, true);
                signedXml.ComputeSignature();
            }
            else
            {
                signedXml.ComputeSignature(((SymmetricSecurityKey)this.signingCredentials.SigningKey).GetKeyedHashAlgorithm(this.signingCredentials.SignatureAlgorithm));
            }

            XmlElement signatureElement = signedXml.GetXml();

            MemoryStream ms = new MemoryStream();
            XmlDictionaryWriter keyIdentifierWriter = XmlDictionaryWriter.CreateTextWriter(ms);
            tokenSerializer.WriteKeyIdentifier(keyIdentifierWriter, this.signingCredentials.SigningKeyIdentifier);
            keyIdentifierWriter.Flush();
            ms.Position = 0;

            XmlDocument dom = signatureElement.OwnerDocument;
            dom.Load(ms);

            // Insert the Key Identifier as the last child element.
            signatureElement.AppendChild(dom.DocumentElement.CloneNode(true));

            keyIdentifierWriter.Close();

            return signatureElement;
        }

    }
}
